'''
数据集预处理脚本farumROP
分类数据集
'''
import sys, os
from data_preprocess.utils.crop import crop_resize_save
from data_preprocess.utils.stat import save_classification_stats
import json
import pandas as pd
from tqdm import tqdm
import argparse

def diagnosis_code_map(diagnosis_code):
    """
    将诊断代码映射为文本描述：
      1 -> "Normal"
      2 -> "Retinopathy of Prematurity Pre-Plus"
      3 -> "Retinopathy of Prematurity Plus"
    若诊断代码无效，则直接 raise ValueError。
    """
    mapping = {
        '1': "Normal",
        '2': "Retinopathy of Prematurity Pre-Plus",
        '3': "Retinopathy of Prematurity Plus"
    }
    code_str = str(diagnosis_code).strip()
    if code_str not in mapping:
        raise ValueError(f"Invalid diagnosis code: {diagnosis_code}")
    return mapping[code_str]

def paser_xlsx(label_path, info_path):
    """ 
    解析 Excel 文件，建立以下两个字典：
      1. patient_info_dict: 从 info_path 中读取，键为 Patient.id，
         值为包含 Patient.BirthWeight, Patient.Gestation Age, Patient.Gender 的字典。
      2. label_dict: 从 label_path 中读取，文件中仅有两列有意义：image_name 和 label，
         注意这里 image_name 不包含 .jpg 后缀。建立 image_name -> diagnosis_code 的映射。
         
    参数：
      label_path (str): 标签 Excel 文件路径
      info_path (str): 患者信息 Excel 文件路径

    返回：
      tuple: (patient_info_dict, label_dict)
    """
    # 解析患者信息
    df_info = pd.read_excel(info_path)
    patient_info_dict = {}
    for _, row in df_info.iterrows():
        patient_id = str(row['Patient.id']).strip()
        patient_info_dict[patient_id] = {
            'BirthWeight': row['Patient.BirthWeight'],
            'Gestation Age': row['Patient.Gestation Age'],
            'Gender': row['Patient.Gender']
        }
    
    # 解析标签，假设 Excel 中列名为 image_name 和 label
    df_label = pd.read_excel(label_path,header=1)
    label_dict = {}
    for _, row in df_label.iterrows():
        # image_name 不含扩展名
        image_name = str(row['Image name']).strip() 
        diagnosis = str(row['Label']).strip()
        if diagnosis == "" or diagnosis is None:
            raise ValueError(f"Empty diagnosis for image {image_name} in {label_path}")
        label_dict[image_name] = diagnosis
    return patient_info_dict, label_dict

def gether_data(data_path, tar_path, prefix='farumROP'):
    """
    对数据集进行预处理：
      1. 在 tar_path 下创建目标目录，包含 images/ 子目录。
      2. 读取 data_path/images 下的所有患者目录，每个患者目录内包含若干图像。
      3. 解析 Excel 文件 Dataset_Labels.xlsx 和 Dataset_Details.xlsx，
         得到 label_dict 和 patient_info_dict。
      4. 对每张图像：
           - 利用 os.path.splitext 提取基础文件名（注意文件名中可能包含 '.'），
           - 根据 label_dict 获取诊断代码，若未找到则直接 raise ValueError，
           - 利用 diagnosis_code_map 转换为文本描述，
           - 调用 crop_resize_save 对图像进行裁剪与 resize（目标为 512×512），
             并统一保存为 PNG 格式，保存路径为 tar_path/images 下。
           - 建立标注信息，记录 image_name、image_path、patient_info 以及 diagnosis。
      5. 将生成的标注信息保存为 JSON 文件（annotations.json）。
      
    参数：
      data_path (str): 原始数据集根目录，要求包含 images/ 目录和两个 Excel 文件：
                        Dataset_Labels.xlsx, Dataset_Details.xlsx
      tar_path (str): 预处理后数据存放目录
      prefix (str): 新图像文件名前缀，默认为 "farumROP"
      
    返回：
      dict: 生成的标注信息字典
    """
    # 创建目标目录及 images 子目录
    os.makedirs(tar_path, exist_ok=True)
    images_tar_dir = os.path.join(tar_path, 'images')
    os.makedirs(images_tar_dir, exist_ok=True)
    
    # 解析 Excel 文件
    patient_info_dict, label_dict = paser_xlsx(
        label_path=os.path.join(data_path, 'Dataset_Labels.xlsx'),
        info_path=os.path.join(data_path, 'Dataset_Details.xlsx')
    )
    
    data_dict = {}
    # 遍历每个患者目录
    patients_dir = os.path.join(data_path, 'images')
    patient_list = sorted(os.listdir(patients_dir))
    
    for patient_id in tqdm(patient_list,desc='patient process',unit='patient'):
        if patient_id not in patient_info_dict:
            raise ValueError(f"Patient id '{patient_id}' not found in info file.")
        patient_info = patient_info_dict[patient_id]
        patient_image_dir = os.path.join(patients_dir, patient_id)
        image_name_list = sorted(os.listdir(patient_image_dir))
        
        for image_name in image_name_list:
            # 使用 os.path.splitext 处理文件名，确保正确获取不含扩展名的部分
            base_name, ext = os.path.splitext(image_name)
            new_image_name = f"{prefix}_{base_name}.png"  # 统一保存为 PNG 格式
            tar_image_rel_path = os.path.join('images', new_image_name)
            src_image_path = os.path.join(patient_image_dir, image_name)
            dest_image_path = os.path.join(tar_path, tar_image_rel_path)
            
            # 调用预处理函数：裁剪+resize
            crop_info=crop_resize_save(
                image_path=src_image_path,
                save_path=dest_image_path,
                resize=(512, 512),
                crop_threshold=15
            )
            
            # 根据 base_name 获取诊断代码，注意 Excel 中的 image_name 没有扩展名
            diag_code = label_dict.get(base_name)
            if diag_code is None:
                raise ValueError(f"Image '{base_name}' not found in label file.")
            diagnosis_text = diagnosis_code_map(diag_code)
            if diagnosis_text == "":
                raise ValueError(f"Empty diagnosis text for image '{base_name}' with code '{diag_code}'")
            data_dict[new_image_name] = {
                'image_name': new_image_name,
                'image_path': dest_image_path,
                'original_path':src_image_path,
                'crop_info':crop_info,
                "patient_info": patient_info,
                'diagnosis': {
                    'classification': {
                        'text': diagnosis_text
                    }
                }
            }
    # 统计分类数据分布
    save_classification_stats(data_dict, './experiments/stat/farumrop.json')
    # 保存标注信息到 JSON 文件
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(data_dict, f, indent=4)
    
    return data_dict

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="新数据集预处理程序 for farumROP")
    parser.add_argument("--data_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_dataset/FARUM-ROP",
                        help="原始数据集根目录，要求包含 images/、Dataset_Labels.xlsx 和 Dataset_Details.xlsx")
    parser.add_argument("--tar_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_processed/FARUM-ROP",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="farumROP",
                        help="处理后图片名称前缀，默认为 'farumROP'")
    args = parser.parse_args()
    
    annotations = gether_data(args.data_path, args.tar_path, prefix=args.prefix)
    print("Preprocessing completed. Annotations saved.")
